# -*- coding: utf-8 -*-
from amber import prmtop
import numpy as np
import math
class fragment():
  def __init__(self,amberAtomName,elementName,coordinate,head,tail,name):
    self.amberAtomName = amberAtomName
    self.elementName = elementName
    self.coordinate = coordinate
    self.head = head
    self.tail = tail
    self.name = name
#更新分子坐标
  def updateCoordinate(self,coordinate):
    self.coordinate = coordinate
  
class protein(fragment):
  def __init__(self,amberAtomName,elementName,residueName,residueIndex,coordinate,head,tail):
    self.amberAtomName = amberAtomName
    self.elementName = elementName
    self.residueIndex = residueIndex
    self.residueName = residueName
    self.coordinate = coordinate
    self.head = head
    self.tail = tail

#切肽键,number是残基号,返回一个完整的残基。
  def cutResidue(self,number):
    head = self.residueIndex[number-1][0]
    tail = self.residueIndex[number-1][1]
    name = self.residueName[number-1]
    aminoElement = self.elementName[head-1:tail]
    aminoAmberAtomName = self.amberAtomName[head-1:tail]
    aminoCoordinate = self.coordinate[head-1:tail]
    residue = fragment(aminoAmberAtomName,aminoElement,aminoCoordinate,head,tail,name)
    return residue

#小分子类，因为不需要分块操作，所以不必要给出amber的原子名。
class micromolecule(fragment):
  mass = {'H':1.008,'C':12.011,'N':14.007,'O':15.999,'NA':22.990,'S':32.065,'CL':35.453,\
          'P':30.974}
  def __init__(self,elementName,coordinate,head,tail,name='xx'):
    self.name = name
    self.elementName = elementName
    self.coordinate = coordinate
    self.head = head
    self.tail = tail

  def closestResidue(self,protein):
    aminoLength = 7;interactionDis=4
    massCenter = np.average(self.coordinate,axis=0)
    residueAllNumber = len(protein.residueIndex)
    radii = 0
    t = 0
    for coord in self.coordinate:
      r2 = np.sum((massCenter - coord)**2)
      r = math.sqrt(r2)
      if radii < r:
        d = t
        radii = r
      t += 1
    closestDis = radii+interactionDis
    maxDis = closestDis + aminoLength
    closest = []
    print 'center=',massCenter
    print 'radii= ',radii
    print 'close= ',closestDis
    print 'max=',maxDis
    i = 1
    while i <= residueAllNumber:
      head = protein.residueIndex[i-1][0]
      tail = protein.residueIndex[i-1][1]
      for number in xrange(head-1,tail):
        r2 = np.sum((massCenter - protein.coordinate[number])**2)
        r = math.sqrt(r2)
        if r > maxDis:
          break
        elif r > closestDis:
          continue
        else:
          closest.append(i)
          break
      i += 1
    return closest

      
    
    
    
